package cn.iocoder.boot.redis;

import java.util.BitSet;
import java.util.Random;

public class CuckooFilter {
    private static final int MAX_KICKS = 500;
    private BitSet[] buckets;
    private int numBuckets;
    private int bucketSize;
    private int numItems;

    public CuckooFilter(int numBuckets, int bucketSize) {
        this.numBuckets = numBuckets;
        this.bucketSize = bucketSize;
        this.buckets = new BitSet[numBuckets];
        for (int i = 0; i < numBuckets; i++) {
            buckets[i] = new BitSet(bucketSize);
        }
        this.numItems = 0;
    }

    public boolean contains(String item) {
        int fingerprint = getFingerprint(item);
        int bucket1 = getBucket(item);
        int bucket2 = getAltBucket(bucket1, fingerprint);
        return buckets[bucket1].get(fingerprint) || buckets[bucket2].get(fingerprint);
    }

    public void insert(String item) {
        if (contains(item)) {
            return;
        }
        int fingerprint = getFingerprint(item);
        int bucket1 = getBucket(item);
        int bucket2 = getAltBucket(bucket1, fingerprint);
        if (buckets[bucket1].cardinality() < bucketSize) {
            buckets[bucket1].set(fingerprint);
            numItems++;
        } else if (buckets[bucket2].cardinality() < bucketSize) {
            buckets[bucket2].set(fingerprint);
            numItems++;
        } else {
            Random random = new Random();
            int bucket = random.nextBoolean() ? bucket1 : bucket2;
            int i = 0;
            while (i < MAX_KICKS) {
                int evictedFingerprint = random.nextInt(bucketSize);
                if (!buckets[bucket].get(evictedFingerprint)) {
                    buckets[bucket].set(evictedFingerprint);
                    String evictedItem = getItem(bucket, evictedFingerprint);
                    insert(evictedItem);
                    return;
                }
                i++;
            }
            rehash();
            insert(item);
        }
    }

    private int getFingerprint(String item) {
        return item.hashCode();
    }

    private int getBucket(String item) {
        return Math.abs(item.hashCode()) % numBuckets;
    }

    private int getAltBucket(int bucket, int fingerprint) {
        return bucket ^ (fingerprint % numBuckets);
    }

    private String getItem(int bucketIndex, int fingerprint) {
        return bucketIndex + ":" + fingerprint;
    }

    private void rehash() {
        int newNumBuckets = numBuckets * 2;
        BitSet[] newBuckets = new BitSet[newNumBuckets];
        for (int i = 0; i < newNumBuckets; i++) {
            newBuckets[i] = new BitSet(bucketSize);
        }
        // 在 rehash 方法中
        for (int i = 0; i < numBuckets; i++) {
            BitSet bucket = buckets[i];
            for (int j = 0; j < bucketSize; j++) {
                if (bucket.get(j)) {
                    String item = getItem(i, j); // 使用桶的索引来构造元素
                    int newBucket = getBucket(item);
                    newBuckets[newBucket].set(getFingerprint(item));
                }
            }
        }
        buckets = newBuckets;
        numBuckets = newNumBuckets;
    }
}